package main

import (
	"context"
	"fmt"
	"net"
	"strconv"
	"strings"

	"gitee.com/hgslb/base"
	"gitee.com/hgslb/facade/fdns"
	"gitee.com/hgslb/infra"
	"github.com/miekg/dns"
)

type dnsServer struct {
	GSLBServer
	dnsSrv *dns.Server
}

func (srv *dnsServer) GetName() string {
	return "dns-server"
}

func (srv *dnsServer) Start(ctx *context.Context) error {
	dnsAddr := ":53"
	sever := dns.Server{Addr: dnsAddr, Net: "udp"}
	sever.Handler = dns.HandlerFunc(func(w dns.ResponseWriter, in *dns.Msg) {
		doDnsHandler(w, in)
	})
	srv.dnsSrv = &sever
	go func() {
		base.GslbLog.Infof("Starting DNS server at %s", dnsAddr)
		if err := sever.ListenAndServe(); err != nil {
			base.GslbLog.Infof("Failed to start DNS server:%v", err)
		}
	}()
	return nil
}

func (srv *dnsServer) Stop(ctx *context.Context) {
	base.GslbLog.Infof("Stopping DNS server")
	if srv.dnsSrv != nil {
		srv.dnsSrv.Shutdown()
	}
}

func doDnsHandler(w dns.ResponseWriter, in *dns.Msg) {
	raddr := base.SplitIp(w.RemoteAddr().String())
	country := doIpContry(raddr)
	out := new(dns.Msg)
	out.SetReply(in)
	out.Authoritative = true
	for _, q := range in.Question {
		if doGetDomainZoneIp(country, &q, out) {
			continue
		}
		if doGetLocalDomain(&q, out) {
			continue
		}
		doGetRemote(raddr, &q, out)
	}
	w.WriteMsg(out)
}

func doGetDomainZoneIp(contry string, q *dns.Question, out *dns.Msg) bool {
	if q.Qtype != dns.TypeA && q.Qtype != dns.TypeAAAA {
		return false
	}

	cz, err := infra.GetDomainZoneCounty(q.Name)
	if err != nil {
		return false
	}

	zone, ok := cz.ContryZone[contry]
	if !ok {
		return false
	}

	zoneIps, err := infra.GetDomainZoneIps(q.Name)
	if err != nil {
		return false
	}

	dnsIps, ok := zoneIps.ZoneIps[zone]
	if !ok {
		return false
	}

	count := 0
	for _, ip := range dnsIps.IpList {
		if ip.Type != int32(dns.TypeA) && ip.Type != int32(dns.TypeAAAA) {
			continue
		}
		if ip.Type == int32(dns.TypeA) {
			ans := dns.A{Hdr: dns.RR_Header{Name: q.Name, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 300}, A: net.ParseIP(ip.Ip).To4()}
			out.Answer = append(out.Answer, &ans)
		} else {
			ans := dns.AAAA{Hdr: dns.RR_Header{Name: q.Name, Rrtype: dns.TypeAAAA, Class: dns.ClassINET, Ttl: 300}, AAAA: net.ParseIP(ip.Ip).To16()}
			out.Answer = append(out.Answer, &ans)
		}
		if int32(q.Qtype) == ip.Type {
			count++
		}
	}
	return (count > 0)

}

func doGetLocalDomain(q *dns.Question, out *dns.Msg) bool {
	if q.Qtype != dns.TypeA && q.Qtype != dns.TypeAAAA {
		return false
	}
	dnsIps, err := infra.GetDomain(q.Name)
	if err != nil {
		return false
	}

	count := 0
	for _, ip := range dnsIps.IpList {
		if ip.Type != int32(dns.TypeA) && ip.Type != int32(dns.TypeAAAA) {
			continue
		}
		if int32(q.Qtype) == ip.Type {
			count++
		}
		if ip.Type == int32(dns.TypeA) {
			ans := dns.A{Hdr: dns.RR_Header{Name: q.Name, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 300}, A: net.ParseIP(ip.Ip).To4()}
			out.Answer = append(out.Answer, &ans)
		} else {
			ans := dns.AAAA{Hdr: dns.RR_Header{Name: q.Name, Rrtype: dns.TypeAAAA, Class: dns.ClassINET, Ttl: 300}, AAAA: net.ParseIP(ip.Ip).To16()}
			out.Answer = append(out.Answer, &ans)
		}
	}
	base.GslbLog.Debugf("Type:%d, count:%d", q.Qtype, count)
	return (count > 0)
}

func doGetRemote(raddr string, q *dns.Question, out *dns.Msg) {
	doGetHttps(raddr, q, out)
}

func doGetHttps(raddr string, q *dns.Question, out *dns.Msg) {
	var sub string
	if !isPrivateIP(raddr) {
		sub = fmt.Sprintf("%s/24", raddr)
	}
	rout, err := infra.DoGetAliyun(strconv.Itoa(int(q.Qtype)), q.Name, sub)
	if err != nil {
		base.GslbLog.Warnf("Get {%s} from aliyun is failed", raddr)
		doGetDns(q, out)
		return
	}

	if rout.Answer == nil {
		base.GslbLog.Warnf("Get {%s} from aliyun is null", raddr)
		return
	}

	var ttl int32 = 0
	ipList := make([]*fdns.GlbsIP, 0)
	count := 0
	for _, ans := range rout.Answer {
		if ans.Type != int(dns.TypeA) && ans.Type != int(dns.TypeAAAA) {
			continue
		}
		ip := &fdns.GlbsIP{Ip: ans.Data, Type: int32(ans.Type)}
		ipList = append(ipList, ip)
		if ttl < int32(ans.TTL) {
			ttl = int32(ans.TTL)
		}
		if ans.Type == int(dns.TypeA) {
			it := dns.A{Hdr: dns.RR_Header{Name: q.Name, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 300}, A: net.ParseIP(ans.Data).To4()}
			out.Answer = append(out.Answer, &it)
		} else {
			it := dns.AAAA{Hdr: dns.RR_Header{Name: q.Name, Rrtype: dns.TypeAAAA, Class: dns.ClassINET, Ttl: 300}, AAAA: net.ParseIP(ans.Data).To16()}
			out.Answer = append(out.Answer, &it)
		}
		count++
	}
	if count > 0 {
		base.GslbLog.Infof("{%s} update dns count {%d}", q.Name, count)
		infra.SetDomain(q.Name, ttl, ipList)
	}
}
func doGetDns(q *dns.Question, out *dns.Msg) {
	m := new(dns.Msg)
	m.SetQuestion(dns.Fqdn(q.Name), q.Qtype)
	m.RecursionDesired = true
	con := base.GslbConf.Dns.Forward
	if q.Qtype == dns.TypeAAAA {
		con = base.GslbConf.Dns.Ipv6
	}
	rout, err := dns.Exchange(m, con)
	if err != nil {
		base.GslbLog.Infof("Get %s from remote failed:%s", q.Name, err)
		return
	}

	var ttl int32 = 0
	ipList := make([]*fdns.GlbsIP, 0)
	count := 0
	for _, ans := range rout.Answer {
		var val string
		switch t := ans.(type) {
		case *dns.A:
			val = t.A.String()
		case *dns.AAAA:
			val = t.AAAA.String()
		default:
			continue
		}
		ip := &fdns.GlbsIP{Ip: val, Type: int32(q.Qtype)}
		ipList = append(ipList, ip)
		if ttl < int32(ans.Header().Ttl) {
			ttl = int32(ans.Header().Ttl)
		}
		out.Answer = append(out.Answer, ans)
		count++
	}
	if count > 0 {
		base.GslbLog.Infof("{%s} update dns count {%d}", q.Name, count)
		infra.SetDomain(q.Name, ttl, ipList)
	}
}

func doIpContry(ip string) string {
	if isPrivateIP(ip) {
		return base.GslbConf.Dns.Local
	}
	country, err := base.GslbGeoDb.Country(ip)
	if err != nil {
		base.GslbLog.Infof("Failed to get %s country:%s", ip, err)
		return base.GslbConf.Dns.Local
	}
	base.GslbLog.Infof("Remote:%s, country:%s", ip, country)
	return country
}

func isPrivateIP(addr string) bool {
	if strings.HasPrefix(addr, "127.0") {
		return true
	}
	ip := net.ParseIP(addr)
	return ip.IsPrivate()
}
